#include "TGraph.h"
#include "mpi.h"
#include <math.h>

#define eps 0.0000001
#define row 0
#define col 0

void floydMPI(int n, int* d)
{
	int rank, size;
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);

	int blockAmount = (int)sqrt(double(size));

	int blockSize = n/blockAmount;
	if (n % blockAmount != 0)
		blockSize++;

	int * blockBuf = new int[blockSize * blockSize];
	int * kCol = new int[blockSize];
	int * kRow = new int[blockSize];
	int tempBufSize = blockSize * blockSize + 2 * blockSize;
	int * tempBuf = new int[tempBufSize];
	int position = 0;

	MPI_Status status;

	for (int k = 0; k < n; k++)
	{
		// first time
		if (k == 0)
		{
			if (rank == 0)
			{

				// fill kCol kRow blockBuf for each and send
				for (int p = 1; p < size; p++)
				{
					// fill blockBuf
					for (int i = 0; i < blockSize; i++)
						for (int j = 0; j < blockSize; j++)
							blockBuf[i * blockSize + j] = d[n * (blockSize * (p / blockAmount) + i) + j + blockSize * (p % blockAmount)];

					// fill k

					for (int j = 0; j < blockSize; j++)
					{
						kCol[j] = d[n * (blockSize * (p / blockAmount) + j) + k];
						kRow[j] = d[n * k + j + blockSize * (p % blockAmount)];
					}

					position = 0;

					MPI_Pack(blockBuf, blockSize * blockSize, MPI_INT, tempBuf, tempBufSize * sizeof(int), &position, MPI_COMM_WORLD);
					MPI_Pack(kCol, blockSize, MPI_INT, tempBuf, tempBufSize * sizeof(int), &position, MPI_COMM_WORLD);
					MPI_Pack(kRow, blockSize, MPI_INT, tempBuf, tempBufSize * sizeof(int), &position, MPI_COMM_WORLD);

					MPI_Send(tempBuf, tempBufSize, MPI_INT, p, 0, MPI_COMM_WORLD); 
				}

				// fill blockBuf for 0
				for (int i = 0; i < blockSize; i++)
					for (int j = 0; j < blockSize; j++)
						blockBuf[i * blockSize + j] = d[n * i + j];

				// fill k for 0

				for (int j = 0; j < blockSize; j++)
				{
					kCol[j] = d[n * j + k];
					kRow[j] = d[n * k + j];
				}

			}
			else
			{
				MPI_Recv(tempBuf, tempBufSize, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);

				position = 0;
				MPI_Unpack(tempBuf, tempBufSize * sizeof(int), &position, blockBuf, blockSize * blockSize, MPI_INT, MPI_COMM_WORLD);
				MPI_Unpack(tempBuf, tempBufSize * sizeof(int), &position, kCol, blockSize, MPI_INT, MPI_COMM_WORLD); 
				MPI_Unpack(tempBuf, tempBufSize * sizeof(int), &position, kRow, blockSize, MPI_INT, MPI_COMM_WORLD); 
			}
			delete[] tempBuf;
		}

		for (int i = 0; i < blockSize; i++)
			for (int j = 0; j < blockSize; j++)
				if (kCol[i] >= 0 && kRow[j] >= 0)
					if ((blockBuf[i * blockSize + j] > (kCol[i] + kRow[j])) || ( blockBuf[i * blockSize + j] < 0) )
						blockBuf[i * blockSize +j] =  kCol[i] + kRow[j];

		MPI_Barrier(MPI_COMM_WORLD);


		// send k if needed 
		if ((rank % blockAmount) == k / blockSize)
		{
			for (int j = 0; j < blockSize; j++)
			{
				kCol[j] = blockBuf[blockSize * j + k % blockSize];
			}
			// send to row
			for (int i = 0; i < blockAmount; i++)
			{
				if (blockAmount * (rank / blockAmount) + i != rank)
				{
					MPI_Send(kCol, blockSize, MPI_INT,  blockAmount * (rank / blockAmount) + i, col, MPI_COMM_WORLD); 
					fflush(stdout);
				}
			}
		}
		else
		{
			// recv kCol
			/*for (int p = 0; p < size; p++)
				if ((p % blockAmount == k / blockSize) && (p / blockAmount == rank / blockAmount))
				{
					MPI_Recv(kCol, blockSize, MPI_INT, p, col, MPI_COMM_WORLD, &status);
					fflush(stdout);
				}*/
			MPI_Recv(kCol, blockSize, MPI_INT, blockAmount * (rank / blockAmount) + k / blockSize, col, MPI_COMM_WORLD, &status);

		}


		if ((rank / blockAmount) == k / blockSize)
		{
			for (int j = 0; j < blockSize; j++)
			{
				kRow[j] = blockBuf[blockSize * (k % blockSize) + j];
			}
			// send to column
			for (int i = 0; i < blockAmount; i++)
			{
				if (i * blockAmount + (rank % blockAmount) != rank)
				{
					MPI_Send(kRow, blockSize, MPI_INT, i * blockAmount + rank % blockAmount, row, MPI_COMM_WORLD); 
					fflush(stdout);
				}
			}
		}
		else
		{
			// recv kRow
			for (int p = 0; p < size; p++)
				if ((p / blockAmount == k / blockSize) && (p % blockAmount == rank % blockAmount))
				{
					MPI_Recv(kRow, blockSize, MPI_INT, p, row, MPI_COMM_WORLD, &status);
					fflush(stdout);
				}
		}
		
		if (k == n - 1)
		{
			for (int i = 0; i < blockSize; i++)
				for (int j = 0; j < blockSize; j++)
					if (kCol[i] >= 0 && kRow[j] >= 0)
						if ((blockBuf[i * blockSize + j] > (kCol[i] + kRow[j])) || ( blockBuf[i * blockSize + j] < 0) )
							blockBuf[i * blockSize +j] =  kCol[i] + kRow[j];
		}
	}

	// rewrite d in rank == 0
	if (rank != 0)
		MPI_Send(blockBuf, blockSize * blockSize, MPI_INT, 0, 0, MPI_COMM_WORLD);
	

	if (rank == 0)
	{
		// write own block
		for (int i = 0; i < blockSize; i++)
			for (int j = 0; j < blockSize; j++)
				d[n * i + j] = blockBuf[i * blockSize + j];

		// recv other block
		for (int p = 1; p < size; p++)
		{
			MPI_Recv(blockBuf, blockSize * blockSize, MPI_INT, p, 0, MPI_COMM_WORLD, &status);
			// write block
			for (int i = 0; i < blockSize; i++)
				for (int j = 0; j < blockSize; j++)
					d[n * (blockSize * (p / blockAmount) + i) + j + blockSize * (p % blockAmount)] = blockBuf[i * blockSize + j];
		}

	}
}